"""
两流体子通道分析程序主程序

本模块是程序的入口点，负责：
1. 读取输入参数
2. 初始化计算网格和求解器
3. 控制计算过程
4. 输出计算结果

作者：[您的名字]
日期：[创建日期]
"""

import os
import json
import numpy as np
from typing import Dict, List
import matplotlib.pyplot as plt

from .physics.properties import FluidProperties
from .physics.two_fluid_model import FluidState
from .numerical.discretization import StaggeredGrid, Discretization
from .numerical.simple import SimpleSolver
from .utils.logger import get_logger
from .postprocess.visualization import ResultVisualizer

class TwoFluidSolver:
    """两流体子通道分析程序主类"""
    
    def __init__(self, config_file: str):
        """
        初始化求解器
        
        参数:
            config_file: str, 配置文件路径
        """
        self.logger = get_logger()
        
        # 读取配置文件
        with open(config_file, 'r') as f:
            self.config = json.load(f)
            
        # 验证子通道结构
        self.verify_channel_structure()
            
        # 初始化组件
        self._initialize_components()
        
    def verify_channel_structure(self):
        """验证子通道结构的正确性"""
        try:
            print("\n====== 验证子通道结构 ======")
            
            # 1. 检查子通道数量
            n_channels = self.config['grid']['n_channels']
            if n_channels != 16:
                raise ValueError(f"子通道数量必须为16,当前为{n_channels}")
                
            # 2. 检查子通道类型
            channel_types = self.config['boundary_conditions']['channel_types']
            if len(channel_types) != 16:
                raise ValueError("channel_types必须定义16个子通道")
                
            # 统计各类型子通道数量
            type_count = {
                'corner': 0,
                'edge': 0,
                'center': 0
            }
            for ch_type in channel_types.values():
                type_count[ch_type] += 1
                
            if type_count['corner'] != 4:
                raise ValueError(f"角通道数量必须为4,当前为{type_count['corner']}")
            if type_count['edge'] != 8:
                raise ValueError(f"边通道数量必须为8,当前为{type_count['edge']}")
            if type_count['center'] != 4:
                raise ValueError(f"中心通道数量必须为4,当前为{type_count['center']}")
                
            # 3. 检查连接关系
            connections = self.config['boundary_conditions']['channel_connections']
            if len(connections) != 16:
                raise ValueError("channel_connections必须定义16个子通道的连接")
                
            # 验证连接的对称性
            for i, conn_list in connections.items():
                i = int(i)
                for j in conn_list:
                    if i not in connections[str(j)]:
                        raise ValueError(f"子通道{i}和{j}的连接关系不对称")
                        
            # 验证连接数量
            for i, conn_list in connections.items():
                ch_type = channel_types[str(i)]
                expected_conn = 2 if ch_type == 'corner' else (3 if ch_type == 'edge' else 4)
                if len(conn_list) != expected_conn:
                    raise ValueError(f"子通道{i}({ch_type})的连接数量应为{expected_conn},当前为{len(conn_list)}")
                    
            # 4. 检查几何参数
            geometry = self.config['boundary_conditions']['channel_geometry']
            required_params = ['flow_area', 'wetted_perimeter', 'heated_perimeter']
            for ch_type in ['corner', 'edge', 'center']:
                if ch_type not in geometry:
                    raise ValueError(f"缺少{ch_type}型通道的几何参数")
                for param in required_params:
                    if param not in geometry[ch_type]:
                        raise ValueError(f"缺少{ch_type}型通道的{param}参数")
                        
            print("子通道结构验证通过!")
            
        except Exception as e:
            print(f"子通道结构验证失败: {str(e)}")
            raise
            
    def debug_output(self, time_step):
        """输出调试信息"""
        try:
            print(f"\n====== 时间步 {time_step} 调试信息 ======")
            
            # 1. 检查质量守恒
            inlet_mass_flow = 0
            outlet_mass_flow = 0
            for j in range(self.n_channels):
                channel_type = self.channel_types[j+1]
                flow_area = self.channel_geometry[channel_type]['flow_area']
                
                # 入口质量流量
                rho_in = self.get_fluid_properties(self.pressure[0,j], self.T_l[0,j], 'liquid')['density']
                inlet_mass_flow += rho_in * self.liquid_velocity[0,j] * flow_area
                
                # 出口质量流量
                rho_l_out = self.get_fluid_properties(self.pressure[-1,j], self.T_l[-1,j], 'liquid')['density']
                rho_v_out = self.get_fluid_properties(self.pressure[-1,j], self.T_v[-1,j], 'vapor')['density']
                outlet_mass_flow += (rho_l_out * (1-self.void_fraction[-1,j]) * self.liquid_velocity[-1,j] + 
                                   rho_v_out * self.void_fraction[-1,j] * self.vapor_velocity[-1,j]) * flow_area
            
            mass_error = abs(inlet_mass_flow - outlet_mass_flow) / inlet_mass_flow
            print(f"质量守恒误差: {mass_error:.6f}")
            
            # 2. 检查能量守恒
            total_power_in = 0
            total_power_out = 0
            total_heat_added = 0
            
            for j in range(self.n_channels):
                channel_type = self.channel_types[j+1]
                flow_area = self.channel_geometry[channel_type]['flow_area']
                heated_perimeter = self.channel_geometry[channel_type]['heated_perimeter']
                
                # 入口能量
                h_in = self.h_l[0,j]
                total_power_in += inlet_mass_flow * h_in / self.n_channels
                
                # 出口能量
                h_l_out = self.h_l[-1,j]
                h_v_out = self.h_v[-1,j]
                total_power_out += outlet_mass_flow * (
                    (1-self.void_fraction[-1,j]) * h_l_out + 
                    self.void_fraction[-1,j] * h_v_out
                ) / self.n_channels
                
                # 加热功率
                for i in range(self.n_axial):
                    total_heat_added += self.q_wl[i,j] * self.dz
            
            energy_error = abs(total_power_out - total_power_in - total_heat_added) / total_heat_added
            print(f"能量守恒误差: {energy_error:.6f}")
            
            # 3. 检查各子通道类型的平均值
            for type_name in ['corner', 'edge', 'center']:
                void_sum = 0
                temp_sum = 0
                count = 0
                
                for j in range(self.n_channels):
                    if self.channel_types[j+1] == type_name:
                        void_sum += np.mean(self.void_fraction[:,j])
                        temp_sum += np.mean(self.T_l[:,j])
                        count += 1
                
                if count > 0:
                    print(f"\n{type_name}型子通道平均值:")
                    print(f"平均空泡率: {void_sum/count:.4f}")
                    print(f"平均温度: {temp_sum/count:.2f} K")
            
            # 4. 检查数值稳定性
            max_courant = 0
            for j in range(self.n_channels):
                channel_type = self.channel_types[j+1]
                flow_area = self.channel_geometry[channel_type]['flow_area']
                
                # 计算Courant数
                max_velocity = max(np.max(abs(self.liquid_velocity[:,j])), 
                                 np.max(abs(self.vapor_velocity[:,j])))
                courant = max_velocity * self.dt / self.dz
                max_courant = max(max_courant, courant)
            
            print(f"\n最大Courant数: {max_courant:.4f}")
            if max_courant > 1:
                print("警告：Courant数大于1，可能存在数值不稳定")
                
        except Exception as e:
            print(f"输出调试信息时出错: {str(e)}")
            
    def _initialize_components(self):
        """初始化各个组件"""
        # 初始化物性计算器
        self.properties = FluidProperties()
        
        # 初始化网格
        self.grid = StaggeredGrid(self.config)
        
        # 初始化离散化器
        self.discretization = Discretization(self.grid, self.config)
        
        # 初始化求解器
        self.solver = SimpleSolver(self.grid, self.discretization, self.config)
        
        # 初始化后处理器
        self.visualizer = ResultVisualizer(self.config)
        
    def setup_initial_conditions(self) -> Dict[str, FluidState]:
        """
        设置初始条件
        
        返回:
            Dict[str, FluidState]: 初始流体状态
        """
        # 从配置文件读取初始条件
        initial = self.config['initial_conditions']
        
        # 计算初始物性
        liquid_props = self.properties.get_liquid_properties(
            initial['pressure'],
            initial['liquid_temperature']
        )
        vapor_props = self.properties.get_vapor_properties(
            initial['pressure'],
            initial['vapor_temperature']
        )
        
        # 创建初始状态
        state = {
            'liquid': FluidState(
                alpha=initial['void_fraction'],
                rho=liquid_props['density'],
                u=initial['liquid_velocity'],
                v=0.0,
                w=0.0,
                T=initial['liquid_temperature'],
                h=liquid_props['enthalpy'],
                mu=liquid_props['viscosity'],
                k=liquid_props['conductivity'],
                cp=liquid_props['specific_heat']
            ),
            'vapor': FluidState(
                alpha=1.0 - initial['void_fraction'],
                rho=vapor_props['density'],
                u=initial['vapor_velocity'],
                v=0.0,
                w=0.0,
                T=initial['vapor_temperature'],
                h=vapor_props['enthalpy'],
                mu=vapor_props['viscosity'],
                k=vapor_props['conductivity'],
                cp=vapor_props['specific_heat']
            ),
            'pressure': initial['pressure']
        }
        
        return state
        
    def setup_boundary_conditions(self) -> dict:
        """
        设置边界条件
        
        返回:
            dict: 边界条件
        """
        # 从配置文件读取边界条件
        boundary = self.config['boundary_conditions']
        
        # 设置通道类型和连接关系
        channel_types = {}
        channel_connections = {}
        connection_gaps = {}
        
        # 根据子通道布局设置边界条件
        layout = boundary['subchannel_layout']
        for i, row in enumerate(layout):
            for j, channel_type in enumerate(row):
                idx = i * len(row) + j
                channel_types[idx] = channel_type
                
                # 设置通道连接
                connections = []
                if i > 0:  # 上方通道
                    connections.append((i-1) * len(row) + j)
                if i < len(layout)-1:  # 下��通道
                    connections.append((i+1) * len(row) + j)
                if j > 0:  # 左侧通道
                    connections.append(i * len(row) + j-1)
                if j < len(row)-1:  # 右侧通道
                    connections.append(i * len(row) + j+1)
                    
                channel_connections[idx] = connections
                
                # 设置连接间隙
                for conn in connections:
                    connection_gaps[(idx, conn)] = boundary['gap_width']
        
        # 返回边界条件字典
        return {
            'inlet': {
                'pressure': boundary['inlet_pressure'],
                'temperature': boundary['inlet_temperature'],
                'mass_flux': boundary['inlet_mass_flux']
            },
            'outlet': {
                'pressure': boundary['outlet_pressure']
            },
            'wall': {
                'heat_flux': boundary['wall_heat_flux']
            },
            'channel_types': channel_types,
            'channel_connections': channel_connections,
            'connection_gaps': connection_gaps,
            'channel_geometry': boundary['channel_geometry']
        }
        
    def run(self):
        """运行计算"""
        # 设置初始条件和边界条件
        current_state = self.setup_initial_conditions()
        boundary_conditions = self.setup_boundary_conditions()
        
        # 时间推进
        time = 0.0
        step = 0
        max_time = self.config['simulation']['max_time']
        save_interval = self.config['output']['save_interval']
        
        # 保存初始结果
        self.save_results(step, time, current_state)
        
        while time < max_time:
            self.logger.info(f"Step {step}, Time: {time:.6f} s")
            
            # 检查数值稳定性
            self.check_numerical_stability()
            
            # 保存上一时刻状态
            prev_state = current_state.copy()
            
            # 求解当前时间步
            new_state, converged = self.solver.solve_time_step(
                current_state,
                prev_state,
                boundary_conditions
            )
            
            if not converged:
                self.logger.error("计算未收敛，程序终止")
                break
                
            # 更新状态
            current_state = new_state
            
            # 更新时间
            time += self.solver.dt
            step += 1
            
            # 输出调试信息
            if step % save_interval == 0:
                self.debug_output(step)
                self.save_results(step, time, current_state)
                
            # 检查是否需要调整时间步长
            self._adjust_time_step()
            
    def save_results(self, step: int, time: float, state: Dict[str, FluidState]):
        """保存计算结果"""
        # 创建输出目录
        output_dir = self.config['output']['directory']
        os.makedirs(output_dir, exist_ok=True)
        
        # 保存状态数据
        filename = os.path.join(output_dir, f"state_{step:06d}.npz")
        np.savez(
            filename,
            time=time,
            liquid_alpha=state['liquid'].alpha,
            liquid_rho=state['liquid'].rho,
            liquid_u=state['liquid'].u,
            liquid_v=state['liquid'].v,
            liquid_w=state['liquid'].w,
            liquid_T=state['liquid'].T,
            liquid_h=state['liquid'].h,
            vapor_alpha=state['vapor'].alpha,
            vapor_rho=state['vapor'].rho,
            vapor_u=state['vapor'].u,
            vapor_v=state['vapor'].v,
            vapor_w=state['vapor'].w,
            vapor_T=state['vapor'].T,
            vapor_h=state['vapor'].h,
            pressure=state['pressure']
        )
        
        # 生成可视化结果
        self.visualizer.plot_results(state, step, time)
        
    def _adjust_time_step(self):
        """调整时间步长"""
        # 计算当前的CFL数
        max_courant = 0
        for j in range(self.n_channels):
            channel_type = self.channel_types[j+1]
            flow_area = self.channel_geometry[channel_type]['flow_area']
            
            # 计算最大速度
            max_velocity = max(np.max(abs(self.liquid_velocity[:,j])), 
                             np.max(abs(self.vapor_velocity[:,j])))
            courant = max_velocity * self.dt / self.dz
            max_courant = max(max_courant, courant)
        
        # 基于CFL条件调整时间步长
        target_courant = 0.8  # 目标CFL数
        if max_courant > 0:
            dt_cfl = self.dt * target_courant / max_courant
        else:
            dt_cfl = self.solver.max_time_step
            
        # 基于收敛性能调整
        if self.solver.last_iterations < self.solver.max_iterations * 0.5:
            # 收敛快，可以适当增加时间步长
            dt_conv = self.dt * 1.1
        elif self.solver.last_iterations > self.solver.max_iterations * 0.8:
            # 收敛慢，需要减小时间步长
            dt_conv = self.dt * 0.8
        else:
            dt_conv = self.dt
            
        # 取两种调整方式的最小值
        new_dt = min(dt_cfl, dt_conv)
        
        # 确保时间步长在允许范围内
        new_dt = max(self.solver.min_time_step, 
                    min(self.solver.max_time_step, new_dt))
        
        # 如果时间步长变化超过20%，才进行调整
        if abs(new_dt - self.dt) / self.dt > 0.2:
            self.logger.info(f"调整时间步长: {self.dt:.6f} -> {new_dt:.6f}")
            self.dt = new_dt

def main():
    """程序入口点"""
    import argparse
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='两流体子通道分析程序')
    parser.add_argument('config', help='配置文件路径')
    args = parser.parse_args()
    
    # 创建并运行求解器
    solver = TwoFluidSolver(args.config)
    solver.run()

if __name__ == '__main__':
    main() 